package com.oscar.springaidemo.ai.config.nodeType;

import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.action.NodeAction;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.prompt.PromptTemplate;
import reactor.core.publisher.Flux;

import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;


@Slf4j
/**
 * 节点类型，这里是拓展节点，可以给每个节点对应一个system message
 */
public class ExpaderNode implements NodeAction {

    private static final PromptTemplate A = new PromptTemplate(
            """
                    你是一个在互联网电商领域工作20年的资深程序员，以幽默风趣的风格回答问题。
                    """
    );

    private final ChatClient chatClient;

    private final Integer NUMBER = 4;

    public ExpaderNode(ChatClient.Builder builder){
        this.chatClient = builder.build();
    }

    public ExpaderNode(ChatClient chatClient){
        this.chatClient = chatClient;
    }

    @Override
    public Map<String, Object> apply(OverAllState state) throws Exception {
        log.info("进行节点记录：{}",state.data());
        String query = state.value("query","");
        Integer expanderNumber = state.value("expander_number",this.NUMBER);

        Flux<String> streamResult =  this.chatClient.prompt().user((user)-> user.text(A.getTemplate())).stream().content();
        String result = streamResult.reduce("", (acc, item) -> acc + item).block();

        List<String> queryVariants = Arrays.asList(result.split("\n"));
        HashMap<String, Object> map = new HashMap<>();
        map.put("expander_content",queryVariants);
        return map;
    }
}
